package com.fastjars.business.spring.http.servlet.wrapper;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.servlet.http.HttpServletResponse;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.PrintWriter;
import java.io.UnsupportedEncodingException;

/**
 * HttpServletResponseWrapper
 *
 * @author Brant Liu <br> 邮箱：<br>lbf1988@qq.com <br>日期：<br>2017/12/24
 * @version 1.0.0
 */
public class HttpServletResponseWrapper extends javax.servlet.http.HttpServletResponseWrapper {
    private static final Logger LOGGER = LoggerFactory.getLogger(HttpServletResponseWrapper.class);

    ByteArrayOutputStream output;
    ServletOutputStream filterOutput;
    XssPrintWriter printWriter;

    public HttpServletResponseWrapper(HttpServletResponse response) {
        super(response);
        output = new ByteArrayOutputStream();
        printWriter = new XssPrintWriter(output);
    }

    @Override
    public javax.servlet.ServletOutputStream getOutputStream() throws IOException {
        if (filterOutput == null) {
            filterOutput = new ServletOutputStream(output);
        }
        return filterOutput;
    }

    public String getContent() {
        try {
            return printWriter.getByteArrayOutputStream().toString("UTF-8");
        } catch (UnsupportedEncodingException e) {
            LOGGER.debug(e.getMessage());
            return "UnsupportedEncoding";
        }
    }

    //覆盖getWriter()方法，使用我们自己定义的Writer
    public PrintWriter getWriter() throws IOException {
        return printWriter;
    }

    public byte[] getDataStream() {
        return output.toByteArray();
    }

    private static class XssPrintWriter extends PrintWriter {
        //此即为存放response输入流的对象
        ByteArrayOutputStream myOutput;

        public XssPrintWriter(ByteArrayOutputStream output) {
            super(output);
            myOutput = output;
        }

        public ByteArrayOutputStream getByteArrayOutputStream() {
            return myOutput;
        }
    }
}
